/*
 * Copyright (c) 2021, Peter Abeles. All Rights Reserved.
 *
 * This file is part of BoofCV (http://boofcv.org).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package boofcv.alg.segmentation.ms;

import boofcv.struct.image.GrayS32;
import georegression.struct.point.Point2D_I32;
import org.ddogleg.struct.DogArray;
import org.ddogleg.struct.DogArray_I32;

/**
 * Merges together regions which have modes close to each other and have a similar color.
 *
 * @author Peter Abeles
 */
public class MergeRegionMeanShift extends RegionMergeTree {

	// maximum distance squared in pixels that two nodes can be apart to be merged
	private final int maxSpacialDistanceSq;
	// Maximum Euclidean distance squared two colors can be for them to be considered similar
	private final float maxColorDistanceSq;

	// Search radius in pixels when looking for regions to merge with
	private final int searchRadius;

	/**
	 * Configures MergeRegionMeanShift
	 *
	 * @param maxSpacialDistance The maximum spacial distance (pixels) at which two modes can be for their
	 * regions to be merged together.
	 * @param maxColorDistance The maximum Euclidean distance two colors can be from each other for them to be merged.
	 */
	public MergeRegionMeanShift( int maxSpacialDistance, float maxColorDistance ) {

		this.searchRadius = maxSpacialDistance;
		this.maxSpacialDistanceSq = maxSpacialDistance*maxSpacialDistance;
		this.maxColorDistanceSq = maxColorDistance*maxColorDistance;
	}

	/**
	 * Merges together similar regions which are in close proximity to each other. After merging
	 * most of the input data structures are modified to take in account the  changes.
	 *
	 * @param pixelToRegion (Input/output) Image that specifies the segmentation. Modified.
	 * @param regionMemberCount (Input/output) Number of pixels in each region. Modified.
	 * @param regionColor (Input/output) Color of each region. Modified.
	 * @param modeLocation (Input) Location of each region's mode. Not modified.
	 */
	public void process( GrayS32 pixelToRegion,
						 DogArray_I32 regionMemberCount,
						 DogArray<float[]> regionColor,
						 DogArray<Point2D_I32> modeLocation ) {
		stopRequested = false;
		initializeMerge(regionMemberCount.size);

		markMergeRegions(regionColor, modeLocation, pixelToRegion);
		if (stopRequested)
			return;

		performMerge(pixelToRegion, regionMemberCount);
	}

	/**
	 * Takes the mode of a region and searches the local area around it for other regions. If the region's mode
	 * is also within the local area its color is checked to see if it's similar enough. If the color is similar
	 * enough then the two regions are marked for merger.
	 */
	protected void markMergeRegions( DogArray<float[]> regionColor,
									 DogArray<Point2D_I32> modeLocation,
									 GrayS32 pixelToRegion ) {
		for (int targetId = 0; targetId < modeLocation.size && !stopRequested; targetId++) {

			float[] color = regionColor.get(targetId);
			Point2D_I32 location = modeLocation.get(targetId);

			int x0 = location.x - searchRadius;
			int x1 = location.x + searchRadius + 1;
			int y0 = location.y - searchRadius;
			int y1 = location.y + searchRadius + 1;

			// ensure that all pixels it examines are inside the image
			if (x0 < 0) x0 = 0;
			if (x1 > pixelToRegion.width) x1 = pixelToRegion.width;
			if (y0 < 0) y0 = 0;
			if (y1 > pixelToRegion.height) y1 = pixelToRegion.height;

			// look at the local neighborhood
			for (int y = y0; y < y1; y++) {
				for (int x = x0; x < x1; x++) {
					int candidateId = pixelToRegion.unsafe_get(x, y);

					// see if it is the same region
					if (candidateId == targetId)
						continue;

					// see if the mode is near by
					Point2D_I32 p = modeLocation.get(candidateId);
					if (p.distance2(location) <= maxSpacialDistanceSq) {

						// see if the color is similar
						float[] candidateColor = regionColor.get(candidateId);
						float colorDistance = SegmentMeanShiftSearch.distanceSq(color, candidateColor);

						if (colorDistance <= maxColorDistanceSq) {
							// mark the two regions as merged
							markMerge(targetId, candidateId);
						}
					}
				}
			}
		}
	}
}
